# All content Copyright 2010 Cyrus Omar <cyrus.omar@gmail.com> unless otherwise
# specified.
#
# Contributors:
#     Cyrus Omar <cyrus.omar@gmail.com>
#
# This file is part of, and licensed under the terms of, the atomic-hedgehog
# package.
#
# The atomic-hedgehog package is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# The atomic-hedgehog package is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
# License for more details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with the atomic-hedgehog package. If not, see <http://www.gnu.org/licenses/>.

"""Utility functions related to numerics and the numpy package."""

import random

import numpy
import ahh.util as util
import ahh.util.proxies as proxies

int_size_bytes = numpy.dtype("int").itemsize
"""The size of the int type in bytes."""

int_size_bits = int_size_bytes * 8
"""The size of the int type in bits."""

if not int_size_bits in (32, 64):
    warn("System integer size (ahh.util.np.int_size_bits == %d) is unsupported." % int_size)

float_size_bytes = numpy.dtype("float").itemsize
"""The size of the float type in bytes."""

float_size_bits = float_size_bytes * 8
"""The size of the float type in bits."""

if not float_size_bits in (32, 64):
    warn("System float size (ahh.util.np.float_size_bits == %d) is unsupported." % float_size_bits)

## Matrix stuff
class JaggedMatrix(proxies.Mask):
    def __init__(self, count=None, wrapped=None):
        if count is not None and wrapped is None:
            if hasattr(count, 'count'): count = count.count
            wrapped = [ [] for i in xrange(count) ]

        proxies.Mask.__init__(wrapped)

    @staticmethod
    def presized(lengths, dtype=numpy.int32):
        return JaggedMatrix(len(lengths), [ numpy.empty(length, dtype)
                                            for length in lengths ])

class AdjacencyMatrix(JaggedMatrix):
    def __init__(self, count, dtype=numpy.int32, mat=None):
        JaggedMatrix.__init__(self, count, mat)
        object.__setattr__(self, 'dtype', dtype)

    def connect_randomly(self, p, (src_start, src_end)=(0, None),
                                (target_start, target_end)=(0, None)):
        n = len(p)
        if src_end is None:
            src_end = n
        if target_end is None:
            target_end = n

        src_num = src_end - src_start
        target_num = target_end - target_start

        degree = numpy.random.binomial(target_num, p, src_num)

        for i in xrange(src_num):
            sample = random.sample(xrange(target_start, target_end), degree[i])
            self[src_start + i].extend(sample)

    @property
    def flattened(self):
        flattened = numpy.empty(self.n_edges, self.dtype)
        return self._flatten_common(flattened, 0)

    @property
    def flattened_with_abs_offsets(self):
        count = self.count
        flattened = numpy.empty(count + self.n_edges, self.dtype)
        flattened[0:count] = self.offsets
        return self._flatten_common(flattened, count)

    @property
    def flattened_with_net_offsets(self):
        count = self.count
        flattened = numpy.empty(count + self.n_edges, self.dtype)
        flattened[0:count] = self.offsets + range(0, count)
        return self._flatten_common(flattened, count)

    @property
    def flattened_with_rel_offsets(self):
        count = self.count
        flattened = numpy.empty(count + self.n_edges, self.dtype)
        flattened[0:count] = self.offsets + range(count, 0, -1)
        return self._flatten_common(flattened, count)

    def _flatten_common(self, flattened, start):
        i = start
        for degree, neighbors in zip(self.degrees, self.mat):
            flattened[i:(i+degree)] = neighbors
            i += count
        return flattened

    @property
    def degrees(self):
        return numpy.array((len(x) for x in self.mat), self.dtype)

    @property
    def n_edges(self):
        degrees = self.degrees
        return numpy.sum(degrees)

    @property
    def offsets(self):
        degrees = self.degrees
        return numpy.cumsum(degrees)
